"""
1. 获取数据
2. 创建神经网络
3. 创建优化器，损失计算，准确率
4. 开始训练
5. 模型评估
"""
import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
import time
import numpy as np
import tensorflow as tf
from tensorflow.python import keras


class SignleNN:
    def __init__(self):
        # 获取数据: 训练特征值， 训练目标值， 测试特征值， 测试目标值
        (self.x_train, self.y_train), (self.x_test, self.y_test) = keras.datasets.fashion_mnist.load_data()
        # print(self.x_train.shape)  # (60000, 28, 28)
        # print(self.y_train.shape)  # (60000,)
        # print(self.x_test.shape)  # (10000, 28, 28)
        # print(self.y_test.shape)  # (10000,)

        # 数据归一化处理
        self.x_train = self.x_train / 255.0
        self.x_test = self.x_test / 255.0

        # 创建神经网络
        self.model = keras.Sequential([
            keras.layers.Flatten(input_shape=(28, 28)),
            keras.layers.Dense(128, activation=tf.nn.relu),
            # keras.layers.Dropout(0.5),  # 当模型发生了过拟合现象时才加入Dropout正则化（L2，L1也是）
            keras.layers.Dense(128, activation=tf.nn.relu),
            keras.layers.Dense(128, activation=tf.nn.relu),
            # keras.layers.Dropout(0.5),
            keras.layers.Dense(128, activation=tf.nn.relu),
            keras.layers.Dense(10, activation=tf.nn.softmax)
        ])

    def snn_compile(self):
        """
        创建优化器，损失计算，准确率
        :return:
        """
        self.model.compile(optimizer=keras.optimizers.Adam(),
                           # loss="sparse_categorical_crossentropy",  # 真实值不为0和1时要用这个交叉熵损失
                           loss=keras.losses.sparse_categorical_crossentropy,  # 和上面等价
                           metrics=["accuracy"])

    def snn_fit(self):
        """
        模型训练
        :return:
        # """
        # # 增加callback项
        # check = keras.callbacks.ModelCheckpoint(
        #     # 1.15版本中准确率是acc损失是loss！！！！
        #     './ckpt/singlenn_{epoch:02d}-{acc:.2f}.h5',
        #     monitor='acc',
        #     save_best_only=True,  # 保存当前轮次准确率（如果第一个参数和第二个参宿是loss，则是损失最小）最好的模型
        #     save_weights_only=True,  # 只保存权重
        #     mode='auto',
        #     period=2)  # 保存2的倍数轮次的模型
        # 再fit中增加TensorBoard项
        board = keras.callbacks.TensorBoard("./graph",  # 保存事件文件目录
                                            write_graph=True  # 是否显示图结构
                                            # histogram_freq=1,
                                            # write_images=True,  # 是否显示图片
                                            # write_grads=True  # 是否显示梯度，histogram_freq必须大于0
                                            )

        # batch_size默认是32
        self.model.fit(self.x_train, self.y_train, epochs=1, batch_size=32, callbacks=[board])

    def snn_evaluate(self):
        """
        模型评估
        :return:
        """
        loss, acc = self.model.evaluate(self.x_test, self.y_test)
        print(loss, acc)

    def snn_predict(self):
        """
        模型的预测
        :return:
        """
        # # 加载模型ckpt形式
        # if os.path.exists("./ckpt/checkpoint"):
        #     self.model.load_weights("./ckpt/SignleNN")
        #
        # return self.model.predict(self.x_test)

        # 模型加载h5形式
        if os.path.exists("./ckpt/sig.h5"):
            self.model.load_weights("./ckpt/sig.h5")

        return self.model.predict(self.x_test)


if __name__ == '__main__':
    snn = SignleNN()
    snn.snn_compile()
    # start_t = time.time()
    snn.snn_fit()
    # end_t = time.time()
    # print("用时:", start_t - end_t)
    snn.snn_evaluate()
    #
    # # 模型的保存, 注意保存时的路径名称，加载的时候必须和这个名称一样
    # # ckpt保存形式
    # snn.model.save_weights("./ckpt/SignleNN")
    # 保存为h5形式
    # snn.model.save_weights("./ckpt/sig.h5")

    # print(np.argmax(snn.snn_predict(), axis=1))
